package com.lnn.springBase.utils.image;

import javax.imageio.ImageIO;
import java.awt.*;
import java.awt.image.BufferedImage;
import java.io.File;
import java.util.ArrayList;
import java.util.List;

public class GridDataConverter {

    // 颜色分段定义
    public static class ColorStop {
        public final Color color;
        public final Double min;
        public final Double max;

        /**
         * @param color 颜色值
         * @param min 区间最小值（null表示自动）
         * @param max 区间最大值（null表示自动）
         */
        public ColorStop(Color color, Double min, Double max) {
            this.color = color;
            this.min = min;
            this.max = max;
        }
    }

    // 图例配置
    public static class LegendConfig {
        public int width = 120;
        public int height = 20;
        public int tickCount = 5;
        public boolean showTitle = true;
        public String title = "Value";
        public Font font = new Font("Arial", Font.PLAIN, 10);
        public Color textColor = Color.BLACK;
        public Color defaultValueColor = null; // null表示透明
    }

    // 结果包装类
    public static class ConversionResult {
        public final BufferedImage image;
        public final BufferedImage legend;
        public final double dataMin;
        public final double dataMax;

        public ConversionResult(BufferedImage image, BufferedImage legend,
                                double dataMin, double dataMax) {
            this.image = image;
            this.legend = legend;
            this.dataMin = dataMin;
            this.dataMax = dataMax;
        }
    }

    /**
     * 数据转图像（自动范围）
     */
    public static ConversionResult convert(double[][] data, List<ColorStop> stops,
                                           LegendConfig legendConfig) {
        double[] range = calculateDataRange(data);
        return convert(data, stops, range[0], range[1], legendConfig);
    }

    /**
     * 数据转图像（自定义范围）
     */
    public static ConversionResult convert(double[][] data, List<ColorStop> stops,
                                           Double userMin, Double userMax,
                                           LegendConfig legendConfig) {
        // 计算有效范围
        double[] autoRange = calculateDataRange(data);
        double effectiveMin = (userMin != null) ? userMin : autoRange[0];
        double effectiveMax = (userMax != null) ? userMax : autoRange[1];

        // 解析颜色分段
        List<ResolvedStop> resolvedStops = resolveStops(stops, effectiveMin, effectiveMax);

        // 生成图像
        BufferedImage image = createImage(data, resolvedStops, effectiveMin, effectiveMax, legendConfig);

        // 生成图例
        BufferedImage legend = createLegend(resolvedStops, effectiveMin, effectiveMax, legendConfig);

        return new ConversionResult(image, legend, effectiveMin, effectiveMax);
    }

    /**
     * 图像转数据
     */
    public static double[][] convertBack(BufferedImage image, List<ColorStop> stops,
                                         double dataMin, double dataMax) {
        List<ResolvedStop> resolvedStops = resolveStops(stops, dataMin, dataMax);
        int width = image.getWidth();
        int height = image.getHeight();
        double[][] data = new double[width][height];

        for (int x = 0; x < width; x++) {
            for (int y = 0; y < height; y++) {
                int rgb = image.getRGB(x, y);
                int alpha = (rgb >> 24) & 0xFF;

                if (alpha == 0) {
                    data[x][y] = Double.NaN;
                } else {
                    Color color = new Color(rgb);
                    data[x][y] = findValue(color, resolvedStops, dataMin, dataMax);
                }
            }
        }
        return data;
    }

    // 内部实现方法
    private static class ResolvedStop {
        final Color color;
        final double min;
        final double max;

        ResolvedStop(Color color, double min, double max) {
            this.color = color;
            this.min = min;
            this.max = max;
        }
    }

    private static List<ResolvedStop> resolveStops(List<ColorStop> stops,
                                                   double globalMin, double globalMax) {
        List<ResolvedStop> resolved = new ArrayList<>();
        for (ColorStop stop : stops) {
            double min = (stop.min != null) ? stop.min : globalMin;
            double max = (stop.max != null) ? stop.max : globalMax;
            resolved.add(new ResolvedStop(stop.color, min, max));
        }
        return resolved;
    }

    private static BufferedImage createImage(double[][] data, List<ResolvedStop> stops,
                                             double dataMin, double dataMax,
                                             LegendConfig config) {
        int width = data.length;
        int height = data[0].length;
        BufferedImage image = new BufferedImage(width, height, BufferedImage.TYPE_INT_ARGB);

        for (int x = 0; x < width; x++) {
            for (int y = 0; y < height; y++) {
                double value = data[x][y];

                if (Double.isNaN(value)) {
                    setDefaultColor(image, x, y, config);
                } else {
                    Color color = findColor(value, stops, dataMin, dataMax);
                    image.setRGB(x, y, color.getRGB());
                }
            }
        }
        return image;
    }

    private static void setDefaultColor(BufferedImage image, int x, int y, LegendConfig config) {
        if (config.defaultValueColor != null) {
            image.setRGB(x, y, config.defaultValueColor.getRGB());
        } else {
            image.setRGB(x, y, 0x00000000); // 完全透明
        }
    }

    private static Color findColor(double value, List<ResolvedStop> stops,
                                   double dataMin, double dataMax) {
        // 查找匹配的分段
        for (ResolvedStop stop : stops) {
            if (value >= stop.min && value <= stop.max) {
                return stop.color;
            }
        }

        // 处理边界情况
        if (value < dataMin) return stops.get(0).color;
        if (value > dataMax) return stops.get(stops.size()-1).color;

        // 线性插值
        return interpolateColor(value, stops, dataMin, dataMax);
    }

    private static Color interpolateColor(double value, List<ResolvedStop> stops,
                                          double dataMin, double dataMax) {
        double position = (value - dataMin) / (dataMax - dataMin);
        int steps = stops.size();
        int index = (int) (position * (steps - 1));
        index = Math.max(0, Math.min(index, steps - 2));

        ResolvedStop start = stops.get(index);
        ResolvedStop end = stops.get(index + 1);
        double ratio = (value - start.max) / (end.min - start.max);

        return blend(start.color, end.color, ratio);
    }

    private static Color blend(Color c1, Color c2, double ratio) {
        ratio = Math.max(0.0, Math.min(1.0, ratio));
        int r = (int) (c1.getRed() * (1 - ratio) + c2.getRed() * ratio);
        int g = (int) (c1.getGreen() * (1 - ratio) + c2.getGreen() * ratio);
        int b = (int) (c1.getBlue() * (1 - ratio) + c2.getBlue() * ratio);
        return new Color(r, g, b);
    }

    private static BufferedImage createLegend(List<ResolvedStop> stops,
                                              double dataMin, double dataMax,
                                              LegendConfig config) {
        BufferedImage legend = new BufferedImage(config.width, config.height, BufferedImage.TYPE_INT_ARGB);
        Graphics2D g = legend.createGraphics();
        g.setRenderingHint(RenderingHints.KEY_ANTIALIASING, RenderingHints.VALUE_ANTIALIAS_ON);

        // 绘制颜色条
        int barWidth = config.width - 40;
        for (int i = 0; i < barWidth; i++) {
            double pos = (double)i / barWidth;
            double value = dataMin + pos * (dataMax - dataMin);
            Color color = findColor(value, stops, dataMin, dataMax);
            g.setColor(color);
            g.drawLine(20 + i, 5, 20 + i, 15);
        }

        // 绘制刻度
        g.setColor(config.textColor);
        g.setFont(config.font);
        for (int i = 0; i <= config.tickCount; i++) {
            double value = dataMin + (dataMax - dataMin) * i / config.tickCount;
            int x = 20 + (barWidth * i) / config.tickCount;
            g.drawLine(x, 16, x, 20);
            String label = String.format("%.2f", value);
            g.drawString(label, x - 15, 35);
        }

        g.dispose();
        return legend;
    }

    private static double[] calculateDataRange(double[][] data) {
        double min = Double.POSITIVE_INFINITY;
        double max = Double.NEGATIVE_INFINITY;
        boolean hasValid = false;

        for (double[] row : data) {
            for (double val : row) {
                if (!Double.isNaN(val)) {
                    hasValid = true;
                    min = Math.min(min, val);
                    max = Math.max(max, val);
                }
            }
        }

        if (!hasValid) return new double[]{0.0, 1.0};
        if (min == max) max += 1e-6;
        return new double[]{min, max};
    }

    private static double findValue(Color target, List<ResolvedStop> stops,
                                    double dataMin, double dataMax) {
        // 查找最近颜色对应的值
        double closestDist = Double.MAX_VALUE;
        double bestValue = dataMin;

        int steps = 1000;
        for (int i = 0; i < steps; i++) {
            double value = dataMin + (dataMax - dataMin) * i / steps;
            Color color = findColor(value, stops, dataMin, dataMax);
            double dist = colorDistance(color, target);
            if (dist < closestDist) {
                closestDist = dist;
                bestValue = value;
            }
        }
        return bestValue;
    }

    private static double colorDistance(Color c1, Color c2) {
        int dr = c1.getRed() - c2.getRed();
        int dg = c1.getGreen() - c2.getGreen();
        int db = c1.getBlue() - c2.getBlue();
        return Math.sqrt(dr*dr + dg*dg + db*db);
    }

    public static void main(String[] args) throws Exception {
        // 创建测试数据
        double[][] data = new double[100][100];
        for (int i = 0; i < 100; i++) {
            for (int j = 0; j < 100; j++) {
                data[i][j] = Math.sqrt(i*i + j*j);
                if (Math.random() < 0.1) data[i][j] = Double.NaN;
            }
        }

        // 配置颜色分段
        List<GridDataConverter.ColorStop> stops = new ArrayList<>();
        stops.add(new GridDataConverter.ColorStop(Color.BLUE, null, 50.0));   // <50
        stops.add(new GridDataConverter.ColorStop(Color.GREEN, 50.0, 100.0)); // 50-100
        stops.add(new GridDataConverter.ColorStop(Color.RED, 100.0, null));   // >100

        // 配置图例
        GridDataConverter.LegendConfig config = new GridDataConverter.LegendConfig();
        config.defaultValueColor = Color.GRAY; // NaN显示为灰色
        config.title = "Distance";
        config.width = 200;
        config.tickCount = 4;

        // 转换数据到图像
        GridDataConverter.ConversionResult result = GridDataConverter.convert(
                data, stops, null, 150.0, config);

        // 保存图像
        ImageIO.write(result.image, "PNG", new File("grid.png"));
        ImageIO.write(result.legend, "PNG", new File("legend.png"));

        // 转换回数据
        double[][] restored = GridDataConverter.convertBack(
                result.image, stops, result.dataMin, result.dataMax);

        // 验证数据
        System.out.println("Original[50][50]: " + data[50][50]);
        System.out.println("Restored[50][50]: " + restored[50][50]);
    }
}
